"""
Sinkhorn距离计算模块

主要功能：
1. 实现Wasserstein距离的快速近似计算
2. 用于评估概率分布之间的距离/相似度
3. 支持批处理计算

特点：
- 基于Sinkhorn-Knopp算法
- 可调节的迭代次数和精度
- 支持GPU加速计算
"""

import torch
import torch.nn as nn

# Adapted from https://github.com/gpeyre/SinkhornAutoDiff

class SinkhornDistance(nn.Module):
    def __init__(self, eps=0.01, max_iter=100, reduction='mean',device='cuda:0'):
        """
        初始化Sinkhorn距离计算器
        Args:
            eps: entropic regularization强度，较小的值会得到更精确但计算更慢的结果
            max_iter: Sinkhorn迭代的最大次数
            reduction: 结果归约方式，'mean'或'sum'
            device: 计算设备
        """
        super(SinkhornDistance, self).__init__()
        self.eps = eps
        self.max_iter = max_iter
        self.reduction = reduction
        self.device=device
        
    def forward(self, x, y):
        """
        计算两个分布之间的Sinkhorn距离
        Args:
            x: 第一个分布的logits
            y: 第二个分布的logits
        Returns:
            cost: Sinkhorn距离值
        """
        # 计算代价矩阵
        C = self._cost_matrix(x, y)  # Wasserstein代价函数
        C = C.to(self.device)
        
        # 设置分布权重（均匀分布）
        n_points = x.shape[-2]
        if x.dim() == 2:
            batch_size = 1
        else:
            batch_size = x.shape[0]

        # 初始化边际分布（uniform distribution）
        mu = torch.empty(batch_size, n_points, dtype=torch.float,
                         requires_grad=False).fill_(1.0 / n_points).squeeze()
        nu = torch.empty(batch_size, n_points, dtype=torch.float,
                         requires_grad=False).fill_(1.0 / n_points).squeeze()

        # 初始化对偶变量
        u = torch.zeros_like(mu)
        u = u.to(self.device)
        v = torch.zeros_like(nu)
        v = v.to(self.device)
        
        # Sinkhorn迭代的收敛阈值
        actual_nits = 0
        thresh = 1e-1

        # Sinkhorn迭代求解最优传输计划
        for i in range(self.max_iter):
            u1 = u  # 记录上一次迭代结果
            # 更新u和v（对偶变量）
            u = self.eps * (torch.log(mu + 1e-8).to(self.device) - self.lse(self.M(C, u, v))).to(self.device) + u
            v = self.eps * (torch.log(nu + 1e-8).to(self.device) - self.lse(self.M(C, u, v).transpose(-2, -1))).to(self.device) + v
            
            # 检查收敛性
            err = (u - u1).abs().sum(-1).mean()
            err = err.to(self.device)

            actual_nits += 1
            if err.item() < thresh:
                break

        # 计算最优传输计划
        U, V = u, v
        pi = torch.exp(self.M(C, U, V))  # 传输计划矩阵
        
        # 计算Sinkhorn距离
        cost = torch.sum(pi * C, dim=(-2, -1))

        # 结果归约
        if self.reduction == 'mean':
            cost = cost.mean()
        elif self.reduction == 'sum':
            cost = cost.sum()

        return cost

    def M(self, C, u, v):
        """
        计算修正的代价矩阵
        Args:
            C: 原始代价矩阵
            u, v: 对偶变量
        Returns:
            M: 修正后的代价矩阵，M_{ij} = (-C_{ij} + u_i + v_j) / eps
        """
        return (-C + u.unsqueeze(-1) + v.unsqueeze(-2)) / self.eps

    @staticmethod
    def _cost_matrix(x, y, p=2):
        """
        计算代价矩阵
        Args:
            x, y: 输入分布
            p: 距离的范数（默认为2，即欧氏距离）
        Returns:
            C: 代价矩阵，C_{ij} = ||x_i - y_j||^p
        """
        x_col = x.unsqueeze(-2)
        y_lin = y.unsqueeze(-3)
        C = torch.sum((torch.abs(x_col - y_lin)) ** p, -1)
        return C

    @staticmethod
    def lse(A):
        """
        计算log-sum-exp
        Args:
            A: 输入张量
        Returns:
            result: log(sum(exp(A))) + 1e-6
        注：添加1e-6防止数值不稳定
        """
        result = torch.log(torch.exp(A).sum(-1) + 1e-6)
        return result

    @staticmethod
    def ave(u, u1, tau):
        "Barycenter subroutine, used by kinetic acceleration through extrapolation."
        return tau * u + (1 - tau) * u1